import gradio as gr
import requests
import json
import os

# 定义 API 地址
BASE_URL = "http://106.54.192.233:8080"  # 替换为实际后台地址
GET_ALL_LEVELS_URL = f"{BASE_URL}/getAllLevels"
GET_LEVEL_LIST_URL = f"{BASE_URL}/getLevelList"  # 正确的获取单个 Level 代码接口
UPLOAD_SAMPLE_FILE_URL = f"{BASE_URL}/updateSampleFile"  # 上传文件接口
GET_CODE_URL = f"{BASE_URL}/getCode"

# 保存文件路径
LOCAL_JSONL_PATH = "lulu_exp/sample_levels.jsonl"
PYTHON_JSONL_PATH = "lulu_exp/sample_levels_python.jsonl"

# 获取所有 Levels 信息
def get_all_levels():
    try:
        response = requests.get(GET_ALL_LEVELS_URL)
        if response.status_code == 200:
            levels = response.json()
            level_titles = ",".join([f'''"{level['title']}"''' for level in levels if level['title'] != ""])
            output_str = "\n".join([f"""ID: {level['id']}, Title: "{level['title']}" """ for level in levels])
            return level_titles, output_str
        else:
            return "", f"获取失败，状态码：{response.status_code}, 响应：{response.text}"
    except Exception as e:
        return "", f"获取时发生错误：{e}"

# 获取每个 level 的代码
def get_level_codes(level_titles):
    titles = [title.strip('"').strip() for title in level_titles.split(",") if title.strip()]
    codes = []
    for title in titles:
        payload = {"titles": [title]}
        try:
            response = requests.post(GET_LEVEL_LIST_URL, json=payload)
            if response.status_code == 200:
                levels = response.json()
                if levels and "data" in levels[0]:
                    code_snippet = levels[0]["data"][:20] + " ... " + levels[0]["data"][-20:]
                    codes.append({"title": title, "code_snippet": code_snippet, "full_code": json.loads(levels[0]["data"])})
                else:
                    codes.append({"title": title, "code_snippet": "代码字段缺失", "full_code": ""})
            else:
                codes.append({"title": title, "code_snippet": f"获取失败，状态码：{response.status_code}"})
        except Exception as e:
            codes.append({"title": title, "code_snippet": f"获取时发生错误：{e}"})
    return codes

# 组织 sample_codes.jsonl
def generate_sample_codes(codes):
    sample_lines = [{"query_str": code["title"], "js_code": code["full_code"]} for code in codes if code["full_code"] != ""]
    # return json.dumps(sample_lines, ensure_ascii=False, indent=2)

    ans = ""
    for line in sample_lines:
        ans += json.dumps(line, ensure_ascii=False) + "\n"
    return ans

# 保存到本地
def save_to_local(sample_codes_jsonl):
    os.makedirs(os.path.dirname(LOCAL_JSONL_PATH), exist_ok=True)
    with open(LOCAL_JSONL_PATH, "w", encoding="utf-8") as f:
        f.write(sample_codes_jsonl)
    return f"保存成功：{LOCAL_JSONL_PATH} "

# 上传到服务器
def upload_to_server(sample_codes_jsonl):
    save_to_local(sample_codes_jsonl)  # 确保文件先保存到本地
    try:
        files = {
            "sample_levels": (os.path.basename(LOCAL_JSONL_PATH), open(LOCAL_JSONL_PATH, "rb"), "application/jsonl"),
            "sample_levels_python": (os.path.basename(PYTHON_JSONL_PATH), open(PYTHON_JSONL_PATH, "rb"), "application/jsonl"),
        }
        response = requests.post(UPLOAD_SAMPLE_FILE_URL, files=files)
        if response.status_code == 200:
            return f"上传成功：{response.json()}"
        else:
            return f"上传失败，状态码：{response.status_code}, 响应：{response.text}"
    except Exception as e:
        return f"上传时发生错误：{e}"


# 调用后端接口生成代码
def get_code(query, model_choice):
    payload = {
        "query": query,
        "mode": "js",  # mode 固定为 "js"
        "model_choice": model_choice
    }
    try:
        response = requests.post(GET_CODE_URL, json=payload)
        if response.status_code == 200:
            result = response.json()
            if "error" in result:
                return f"生成失败：{result['error']}"
            return result  # 返回生成的代码
        else:
            return f"请求失败，状态码：{response.status_code}, 响应：{response.text}"
    except Exception as e:
        return f"生成代码时发生错误：{e}"

# 构建 Gradio 界面
with gr.Blocks() as demo:
    with gr.Tab("生成代码"):
        # 第一行: 输入 Query
        with gr.Row():
            gr.Markdown("### 输入 Query")
            query_input = gr.Textbox(label="Query", placeholder="请输入生成代码的描述")

        # 第二行: 选择模型
        with gr.Row():
            gr.Markdown("### 选择模型")
            model_choice = gr.Radio(
                ["Qwen", "DeepSeek", "Zhipu"],
                value="Qwen",  # 默认选项
                label="模型选择"
            )

        # 第三行: 按钮和生成代码展示
        with gr.Row():
            gr.Markdown("### 点击生成代码")
            generate_button = gr.Button("生成代码")
        code_output = gr.TextArea(label="生成的代码", interactive=False)

        # 按钮交互逻辑
        generate_button.click(
            fn=get_code,
            inputs=[query_input, model_choice],
            outputs=code_output
        )

    with gr.Tab("管理RAG"):
        level_codes = gr.State([])

        # 第一行: 获取所有 Levels 信息
        with gr.Row():
            gr.Markdown("### 获取所有 Levels 信息")
            get_all_button = gr.Button("获取所有 Levels")
        levels_text_area = gr.TextArea(label="所有 Levels 信息", interactive=False)

        # 第二行: 显示和修改 Level Titles
        with gr.Row():
            gr.Markdown("### 修改 Level Titles")
            level_titles = gr.TextArea(label="Level Titles", interactive=True)

        # 第三行: 获取代码
        with gr.Row():
            gr.Markdown("### 获取每个 Level 的代码")
            get_codes_button = gr.Button("获取代码")
        codes_text_area = gr.TextArea(label="每个 Level 的代码", interactive=False)

        # 第四行: 组织成 sample_codes
        with gr.Row():
            gr.Markdown("### 生成 sample_codes.jsonl")
            generate_button = gr.Button("生成 sample_codes.jsonl")
        sample_codes_area = gr.TextArea(label="sample_codes.jsonl", interactive=False)

        # 新增行：保存与上传
        with gr.Row():
            save_button = gr.Button("保存到本地")
            upload_button = gr.Button("上传到服务器")
        action_log = gr.Textbox(label="操作日志", interactive=False)

    # 按钮交互逻辑
    get_all_button.click(
        fn=get_all_levels,
        inputs=[],
        outputs=[level_titles, levels_text_area]
    )

    get_codes_button.click(
        fn=get_level_codes,
        inputs=level_titles,
        outputs=level_codes
    ).then(
        fn=lambda codes: "\n".join([f"Title: {c['title']}, Code Snippet: {c['code_snippet']}" for c in codes]),
        inputs=level_codes,
        outputs=codes_text_area
    )

    generate_button.click(
        fn=generate_sample_codes,
        inputs=level_codes,
        outputs=sample_codes_area
    )

    save_button.click(
        fn=save_to_local,
        inputs=sample_codes_area,
        outputs=action_log
    )

    upload_button.click(
        fn=upload_to_server,
        inputs=sample_codes_area,
        outputs=action_log
    )

# 启动 Gradio 界面
demo.launch()
